import os
import SimpleITK as sitk
import nrrd
import numpy as np
import pydicom
import csv
import shutil
"""dicom_flip_0703
    此将本将编写数据翻转逻辑

    主要参考nrrd header中的朝向矩阵以及厂家号进行翻转

    将医生建议ImageOrientationPatient加入到前转换脚本
"""

"""
    翻转的厂家:
        Philips(大部分为反转，偶尔有正的)
"""
def make_manu_csv(root_start: str, csv_path: str):
    root_dcm_dict = {}
    with open(csv_path, 'w', newline='') as f:
        writer = csv.writer(f)
        header = ['AccessionNumber', 'Manufacturer']
        writer.writerow(header)
        for root, dirs, files in os.walk(root_start):
            filter_dcm_file = []

            for each_file in files:
                if each_file.lower().endswith(".dcm"):
                    filter_dcm_file.append(each_file)

            if len(filter_dcm_file) == 0:
                continue

            # TODO: 文件夹可能存在多个脏数据
            reader = sitk.ImageSeriesReader()
            root_dcm_dict[root] = reader.GetGDCMSeriesIDs(root)[-1]
            series_file_names = reader.GetGDCMSeriesFileNames(root, root_dcm_dict[root])
            ds = pydicom.dcmread(series_file_names[0])
            id = ds.AccessionNumber
            manu = ds.Manufacturer
            data = [id, manu]
            writer.writerow(data)


def reorient_nrrd(input_dir: str, save_dir: str, csv_path: str):
    os.makedirs(save_dir, exist_ok=True)
    special_manufacturer = ['philips']
    for file in os.listdir(input_dir):
        if file.endswith('.nrrd'):
            status = 0

            file_path = os.path.join(input_dir, file)
            data, header = nrrd.read(file_path)
            save_path = os.path.join(save_dir, file)

            for key in range(len(header['space directions'][2])):
                if header['space directions'][2][key] < 0 :
                    header['space directions'][2] = header['space directions'][2] * (-1)
                    status = -1
                    break

            if search_csv(file.replace('.nrrd', ''), csv_path).lower() not in special_manufacturer:
                status = 0

            if status == 0:
                shutil.copy(file_path, save_dir)
                print(f"copied {file}")
            else:
                nrrd.write(save_path, data, header)
                print(f"reoriented {file}")


def search_csv(id: str, csv_path: str):
    # 打开并读取CSV文件
    with open(csv_path, mode='r') as file:
        reader = csv.reader(file)

        # 读取第一行（列名）
        headers = next(reader)

        # 确保第一行有两列并且包含 'column_name' 和 'value'
        if len(headers) < 2:
            raise ValueError("CSV file must contain at least two columns")

        # 获取列索引
        column_name_index = headers.index('AccessionNumber')
        value_index = headers.index('Manufacturer')

        # 遍历文件的每一行数据
        for row in reader:
            # 如果当前行的第一列（列名）与搜索列名匹配
            if row[column_name_index] == id:
                return row[value_index]


def flip_niigz(input_nii: str, output_path: str):

    img = sitk.ReadImage(input_nii)
    image_array = sitk.GetArrayFromImage(img)

    # 使用numpy进行上下翻转
    flipped_array = np.flip(image_array, axis=0)  # axis=0 表示对第一个维度（通常是z轴）进行翻转

    # 将翻转后的numpy数组转换回SimpleITK图像
    flipped_image = sitk.GetImageFromArray(flipped_array)
    flipped_image.CopyInformation(img)
    sitk.WriteImage(flipped_image, output_path)



make_manu_csv(r'B:\projects\dicom_data_test',r'B:\projects\manu.csv')
reorient_nrrd(r"B:\projects\dicom_save",r"B:\projects\manu\manu-2", r"B:\projects\manu.csv")
